{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a2731495-f222-4168-95cf-8f94a477ff43",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 分bin对 $\\alpha$ $\\beta$ 和M做一次mcmc\n",
    "a,b,m 同时拟合，-3<a<3,-5<b<5,-24.3<M<-14.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "df5fd112-1fa2-4299-8f6c-a84591ea5964",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "789faff4-273e-4aa3-a5b3-591c8df17c7f",
   "metadata": {},
   "source": [
    "### 读取模拟数据，0517的z在0.1-0.6之间"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "821d9e74-fe07-4178-9bd5-0ee0ebc9e2bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "simu_obs = pd.read_pickle(\"result-5yr-3cadence-0614-300s.analyze.MacOS\", )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55f08d7b-218e-4281-a716-b75d48ba4631",
   "metadata": {},
   "source": [
    "### 数据分bin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "76aeae45-c6d2-4f6a-91d8-44ebda61a1a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_simu(zmin, zmax, simu):\n",
    "    simu_bin = simu.loc[(simu['z'] >= zmin) & (simu['z'] < zmax)]\n",
    "    return simu_bin"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6192666-e725-4075-a699-82359feea5d4",
   "metadata": {},
   "source": [
    "### 分bin用chisq最小来确定每组的M"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "adbf5824-9fe4-4904-84e1-8ebc35841704",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chisq_bin(params, obs, m_fid):\n",
    "    alpha, beta, M = params\n",
    "    # print(f\"{alpha=}, {beta=}, {M=}, {m_obs=}, {m_fid=}\")\n",
    "    value = alpha**2 + (beta-2)**2 + (M+19)**2\n",
    "    return value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "de65f0fa-b7d0-4a3d-848d-5c2e1520c622",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy import cosmology as cosmo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fe5f34ba-fa77-4985-b1a2-0476a3bb01e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = cosmo.Planck18"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7a59ab25-9b45-47a4-94af-9543f7abef44",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chisq_bin(params, obs, zmin, zmax, sig_n = 0.1):\n",
    "    alpha, beta, M = params # to be fitted by mcmc\n",
    "    simu_bin = select_simu(zmin, zmax, obs)\n",
    "    mu_simu = simu_bin.mb + alpha* simu_bin.x1 - beta*simu_bin.c - M\n",
    "    mu_theory = model.distmod(simu_bin.z.to_numpy()).value\n",
    "    variance = (simu_bin.V_mb + alpha**2 * simu_bin.V_x1 + beta**2 * simu_bin.V_c +\n",
    "                2*alpha*simu_bin.V_mb_x1 - 2*beta*simu_bin.V_mb_c - 2*alpha*beta*simu_bin.V_x1_c) + sig_n \n",
    "    #print(variance)\n",
    "    return ((mu_simu - mu_theory)**2/variance).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2fa4c66f-ceaa-4ee7-8bc4-db887a34baf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#chisq_tot\n",
    "def chisq_total(params, obs, z_edges, sig_n):\n",
    "    # convert protential lists to numpy array\n",
    "    params = np.asarray(params)\n",
    "    z_edges = np.asarray(z_edges)\n",
    "    \n",
    "    # split params to specific arguments\n",
    "    alpha = params[0]\n",
    "    beta  = params[1]\n",
    "    M_bins= params[2:]\n",
    "    # alpha, beta, *M_bins = params # Here * means M_bins is an array\n",
    "    z_bin = (z_edges[:-1] + z_edges[1:])/2   # bins的中心坐标\n",
    "    \n",
    "    chisq_total_value = 0\n",
    "    for bin_id in range(len(z_bin)):\n",
    "        zmin = z_edges[bin_id]\n",
    "        zmax = z_edges[bin_id+1]\n",
    "        M = M_bins[bin_id]\n",
    "        chisq_total_value += chisq_bin([alpha, beta, M], obs, zmin, zmax, sig_n)\n",
    "    \n",
    "    return chisq_total_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c88cf4b7-8efd-4cb0-b40a-49cf1dbd65e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_prob(params, obs, z_edges, sig_n):\n",
    "    # convert protential iterables to numpy array\n",
    "    params = np.asarray(params)\n",
    "    alpha = params[0]\n",
    "    beta = params[1]\n",
    "    M_bins = params[2:]\n",
    "    \n",
    "    # alpha, beta, M = param\n",
    "    if np.any((M_bins < -24.3) | (M_bins > -14.3)):\n",
    "        return -np.inf\n",
    "    if (alpha < -3) or (alpha > 3):\n",
    "        return -np.inf #超过这个范围就不算了\n",
    "    if (beta < -5) or (beta > 5):\n",
    "        return -np.inf   \n",
    "    return -0.5 *chisq_total(params, obs, z_edges, sig_n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8326cf37-eb30-44e9-b91f-833d745b7c3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "nbins = 15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4676185e-b581-4d70-a69d-fd06fff7dabf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['alpha',\n",
       " 'beta',\n",
       " 'M0',\n",
       " 'M1',\n",
       " 'M2',\n",
       " 'M3',\n",
       " 'M4',\n",
       " 'M5',\n",
       " 'M6',\n",
       " 'M7',\n",
       " 'M8',\n",
       " 'M9',\n",
       " 'M10',\n",
       " 'M11',\n",
       " 'M12',\n",
       " 'M13',\n",
       " 'M14']"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "param_names = [\"alpha\", \"beta\", ] + [f\"M{i}\" for i in np.arange(nbins)]\n",
    "param_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "af77e494-79ab-495f-8e99-7a8523b46ec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ndim = len(param_names)\n",
    "nwalkers = 128\n",
    "p0 = np.random.rand(nwalkers, ndim)\n",
    "p0[:, 0] = p0[:, 0] * 6 - 3   # alpha\n",
    "p0[:, 1] = p0[:, 1] * 10 - 5    # beta\n",
    "p0[:, 2:] = p0[:, 2:] * 10 - 24.3    # M"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7241a6ee-076e-4da4-bb7a-2b6c504c26c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.03      , 0.13133333, 0.23266667, 0.334     , 0.43533333,\n",
       "       0.53666667, 0.638     , 0.73933333, 0.84066667, 0.942     ,\n",
       "       1.04333333, 1.14466667, 1.246     , 1.34733333, 1.44866667,\n",
       "       1.55      ])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z_edges = np.linspace(0.03, 1.55, nbins+1, endpoint=True)\n",
    "z_edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa726618-a0ef-4150-863e-2caf00ab78fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "simu_obs.good_fit.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1918a573-be6e-4173-b32e-ce7da779cc8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import emcee"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d0abeeeb-e12d-4352-9b21-fb53bfbad186",
   "metadata": {},
   "outputs": [],
   "source": [
    "# MCMC: saving configure\n",
    "\n",
    "mcmc_filename = f\"salt2mu_abmlarge_0614newsimu300s.h5\"\n",
    "\n",
    "backend = emcee.backends.HDFBackend(mcmc_filename) #在mcmc的过程中查看收敛情况\n",
    "backend.reset(nwalkers, ndim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "aebbdfd5-04fa-44f2-a033-c9f014d54391",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sampler = emcee.EnsembleSampler(nwalkers, ndim, log_prob, args=(simu_obs,z_edges,0.1), backend=backend)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f572c33-6bd3-408a-931a-92df0259dd42",
   "metadata": {},
   "source": [
    "state = sampler.run_mcmc(p0, 100)\n",
    "sampler.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52b1fed9-5fac-473b-b710-084420a6a4c8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 2201/100000 [47:46:22<277:49:50, 10.23s/it]    "
     ]
    }
   ],
   "source": [
    "# MCMC: run it! 监测收敛情况\n",
    "max_n = 100000 # max_n x walkers = 采集的样本数\n",
    "\n",
    "# We'll track how the average autocorrelation time estimate changes\n",
    "index = 0\n",
    "autocorr = np.empty(max_n)\n",
    "\n",
    "# This will be useful to testing convergence\n",
    "old_tau = np.inf\n",
    "\n",
    "# Now we'll sample for up to max_n steps\n",
    "for sample in sampler.sample(p0, iterations=max_n, progress=True):\n",
    "    # Only check convergence every 100 steps\n",
    "    if sampler.iteration % 100:\n",
    "        continue\n",
    "\n",
    "    # Compute the autocorrelation time so far\n",
    "    # Using tol=0 means that we'll always get an estimate even\n",
    "    # if it isn't trustworthy\n",
    "    tau = sampler.get_autocorr_time(tol=0)\n",
    "    autocorr[index] = np.mean(tau)\n",
    "    index += 1\n",
    "\n",
    "    # Check convergence\n",
    "    converged = np.all(tau * 100 < sampler.iteration)\n",
    "    converged &= np.all(np.abs(old_tau - tau) / tau < 0.01)\n",
    "    if converged:\n",
    "        break\n",
    "    old_tau = tau"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e756133-995e-49fb-ae09-846289e77253",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
